import tushare as ts
import pandas as pd
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
import logging

logger = logging.getLogger(__name__)

class StockDataCollector:
    def __init__(self, api_token, config):
        # 初始化Tushare API
        self.api_token = api_token
        self.config = config
        self.pool = ThreadPoolExecutor(max_workers=config.get('max_workers', 5))
        
        # 更健壮的API初始化
        try:
            self.pro = ts.pro_api(api_token)
            # 简单测试API连接
            test_result = self.pro.stock_basic(exchange='', list_status='L', fields='ts_code,symbol,name')
            if test_result is None or test_result.empty:
                logger.warning("Tushare API连接成功但未返回数据，可能是Token权限问题")
            else:
                logger.info(f"Tushare API连接成功，获取到{len(test_result)}条股票基本信息")
        except Exception as e:
            logger.error(f"Tushare API初始化失败: {str(e)}")
            self.pro = None  # 标记为初始化失败
        
    def get_realtime_quotes(self, stock_codes):
        """获取实时行情"""
        try:
            # 尝试使用Tushare获取实时数据
            df = self.pro.daily(ts_code=','.join(stock_codes), trade_date=datetime.now().strftime('%Y%m%d'))
            if df.empty:
                return self._get_mock_quotes(stock_codes)
            return df.to_dict('records')
        except Exception as e:
            logger.warning(f"获取实时行情失败，使用模拟数据: {str(e)}")
            return self._get_mock_quotes(stock_codes)
        
    def _get_mock_quotes(self, stock_codes):
        """生成模拟行情数据（当API不可用时使用）"""
        mock_data = []
        for code in stock_codes:
            # 生成随机模拟数据
            base_price = 10.0 + (hash(code) % 100)
            change_pct = (hash(code[::-1]) % 10 - 5) / 100
            
            mock_data.append({
                'ts_code': code,
                'trade_date': datetime.now().strftime('%Y%m%d'),
                'open': round(base_price * (1 - change_pct * 0.5), 2),
                'high': round(base_price * (1 + abs(change_pct)), 2),
                'low': round(base_price * (1 - abs(change_pct)), 2),
                'close': round(base_price, 2),
                'pre_close': round(base_price * (1 - change_pct), 2),
                'change': round(base_price * change_pct, 2),
                'pct_chg': round(change_pct * 100, 2),
                'vol': int(1000000 + (hash(code) % 10000000)),
                'amount': int(base_price * (1000000 + (hash(code) % 10000000))),
                'is_mock': True  # 标记为模拟数据
            })
        return mock_data
    
    def get_daily_data(self, stock_code, start_date, end_date):
        """获取日K线数据"""
        return self.pro.daily(ts_code=stock_code, start_date=start_date, end_date=end_date)
    
    def get_stock_indicators(self, stock_code):
        """获取技术指标数据"""
        # 获取K线数据
        df = self.get_daily_data(stock_code, self.config['start_date'], self.config['end_date'])
        
        # 计算常用技术指标
        df = self._calculate_ma(df)
        df = self._calculate_rsi(df)
        df = self._calculate_macd(df)
        return df
    
    def _calculate_ma(self, df, periods=[5, 10, 20, 30, 60]):
        """计算移动平均线"""
        for period in periods:
            df[f'MA_{period}'] = df['close'].rolling(period).mean()
        return df
    
    def _calculate_rsi(self, df, periods=[6, 12, 24]):
        """计算RSI指标"""
        # 实现RSI计算逻辑
        for period in periods:
            delta = df['close'].diff()
            gain = (delta.where(delta > 0, 0)).rolling(window=period).mean()
            loss = (-delta.where(delta < 0, 0)).rolling(window=period).mean()
            rs = gain / loss
            df[f'RSI_{period}'] = 100 - (100 / (1 + rs))
        return df
        
    def _calculate_macd(self, df, fast=12, slow=26, signal=9):
        """计算MACD指标"""
        # 实现MACD计算逻辑
        exp1 = df['close'].ewm(span=fast, adjust=False).mean()
        exp2 = df['close'].ewm(span=slow, adjust=False).mean()
        df['MACD'] = exp1 - exp2
        df['MACD_signal'] = df['MACD'].ewm(span=signal, adjust=False).mean()
        df['MACD_hist'] = df['MACD'] - df['MACD_signal']
        return df
    
    def try_api_call(self, func_name, **kwargs):
        """通用API调用方法，带有重试和错误处理"""
        if not hasattr(self.pro, func_name):
            logger.error(f"API方法不存在: {func_name}")
            return None
        
        max_retries = 3
        for attempt in range(max_retries):
            try:
                func = getattr(self.pro, func_name)
                result = func(**kwargs)
                return result
            except Exception as e:
                logger.warning(f"API调用失败 ({attempt+1}/{max_retries}): {str(e)}")
                if "token" in str(e).lower():
                    logger.error("Token验证失败，请检查Token是否正确")
                    break  # Token错误不需要重试
                time.sleep(1)  # 等待1秒后重试
            
        return None  # 所有重试都失败